{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autonomous Agent\n",
    "\n",
    "An equity trading simulation to illustrate autonomous agents powered by tools and resources from MCP servers.\n",
    "\n",
    "The simulation has 4 Traders and a Researcher, powered by a slew of MCP servers with tools & resources:\n",
    "\n",
    "1. Our home-made Accounts MCP server (written by our engineering team!)\n",
    "2. Fetch (get webpage via a local headless browser)\n",
    "3. Memory\n",
    "4. Brave Search\n",
    "5. Financial data\n",
    "\n",
    "And a resource to read information about the trader's account, and their investment strategy.\n",
    "\n",
    "The python module, `traders.py` will manage a single trader on our trading floor.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from datetime import datetime\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "from agents import Agent, Runner, trace, Tool\n",
    "from agents.extensions.models.litellm_model import LitellmModel\n",
    "from agents.mcp import MCPServerStdio\n",
    "\n",
    "from accounts_client import read_accounts_resource, read_strategy_resource\n",
    "from accounts import Account\n",
    "\n",
    "load_dotenv(override=True)\n",
    "api_key = os.getenv(\"AMAZON_BEDROCK_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Gathering the MCP params for our trader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "polygon_api_key = os.getenv(\"POLYGON_API_KEY\")\n",
    "polygon_plan = os.getenv(\"POLYGON_PLAN\")\n",
    "\n",
    "is_paid_polygon = polygon_plan == \"paid\"\n",
    "is_realtime_polygon = polygon_plan == \"realtime\"\n",
    "\n",
    "print(is_paid_polygon)\n",
    "print(is_realtime_polygon)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using polygon if paid plan else use `market_server.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if is_paid_polygon or is_realtime_polygon:\n",
    "    market_mcp = {\n",
    "        \"command\": \"uvx\",\n",
    "        \"args\": [\"--from\", \"git+https://github.com/polygon-io/mcp_polygon@master\", \"mcp_polygon\"], \n",
    "        \"env\": {\"POLYGON_API_KEY\": polygon_api_key}\n",
    "    }\n",
    "\n",
    "else:\n",
    "    market_mcp = ({\n",
    "        \"command\": \"uv\",\n",
    "        \"args\": [\"run\", \"market_server.py\"]\n",
    "    })\n",
    "\n",
    "trader_mcp_server_params = [\n",
    "    {\n",
    "        \"command\": \"uv\", \n",
    "        \"args\": [\"run\", \"accounts_server.py\"]\n",
    "    },\n",
    "    {\n",
    "        \"command\": \"uv\", \n",
    "        \"args\": [\"run\", \"push_server.py\"]\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Researcher using fetch mcp server and brave search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "brave_env = {\"BRAVE_API_KEY\": os.getenv(\"BRAVE_API_KEY\")}\n",
    "\n",
    "researcher_mcp_server_params = [\n",
    "    {\n",
    "        \"command\": \"uvx\", \n",
    "        \"args\": [\"mcp-server-fetch\"]\n",
    "    },\n",
    "    {\n",
    "        \"command\": \"npx\", \n",
    "        \"args\": [\"-y\", \"@modelcontextprotocol/server-brave-search\"], \n",
    "        \"env\": brave_env}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating the MCPServerStdio for each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "researcher_mcp_servers = [MCPServerStdio(params, client_session_timeout_seconds=30) for params in researcher_mcp_server_params]\n",
    "trader_mcp_servers = [MCPServerStdio(params, client_session_timeout_seconds=30) for params in trader_mcp_server_params]\n",
    "mcp_servers = trader_mcp_servers + researcher_mcp_servers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Researcher Agent to do market research\n",
    "\n",
    "And turn it into a tool - remember how this works for OpenAI Agents SDK, and the difference with handoffs?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"us.anthropic.claude-3-7-sonnet-20250219-v1:0\"\n",
    "model=LitellmModel(model=model_name, api_key=api_key)\n",
    "\n",
    "instructions = f\"\"\"You are a financial researcher. You are able to search the web for interesting financial news,\n",
    "look for possible trading opportunities, and help with research.\n",
    "Based on the request, you carry out necessary research and respond with your findings.\n",
    "Take time to make multiple searches to get a comprehensive overview, and then summarize your findings.\n",
    "If there isn't a specific request, then just respond with investment opportunities based on searching latest news.\n",
    "The current datetime is {datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_researcher(mcp_servers) -> Agent:\n",
    "    researcher = Agent(\n",
    "        name=\"Researcher\",\n",
    "        instructions=instructions,\n",
    "        model=model,\n",
    "        mcp_servers=mcp_servers,\n",
    "    )\n",
    "    return researcher"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Research agent as a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_researcher_tool(mcp_servers) -> Tool:\n",
    "    researcher = await get_researcher(mcp_servers)\n",
    "    return researcher.as_tool(\n",
    "            tool_name=\"Researcher\",\n",
    "            tool_description=\"This tool researches online for news and opportunities, \\\n",
    "                either based on your specific request to look into a certain stock, \\\n",
    "                or generally for notable financial news and opportunities. \\\n",
    "                Describe what kind of research you're looking for.\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "research_question = \"What's the latest news on Amazon?\"\n",
    "\n",
    "for server in researcher_mcp_servers:\n",
    "    await server.connect()\n",
    "researcher = await get_researcher(researcher_mcp_servers)\n",
    "with trace(\"Researcher\"):\n",
    "    result = await Runner.run(researcher, research_question, max_turns=20)\n",
    "display(Markdown(result.final_output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Look at the trace\n",
    "\n",
    "https://platform.openai.com/traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_strategy = \"You are a day trader that buys and sells shares based on news and market conditions.\"\n",
    "Account.get(\"Ian\").reset(initial_strategy)\n",
    "\n",
    "display(Markdown(await read_accounts_resource(\"Ian\")))\n",
    "display(Markdown(await read_strategy_resource(\"Ian\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### And now - to create our Trader Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_name = \"Ian\"\n",
    "\n",
    "# Using MCP Servers to read resources\n",
    "account_details = await read_accounts_resource(agent_name)\n",
    "strategy = await read_strategy_resource(agent_name)\n",
    "\n",
    "instructions = f\"\"\"\n",
    "You are a trader that manages a portfolio of shares. Your name is {agent_name} and your account is under your name, {agent_name}.\n",
    "You have access to tools that allow you to search the internet for company news, check stock prices, and buy and sell shares.\n",
    "Your investment strategy for your portfolio is:\n",
    "{strategy}\n",
    "Your current holdings and balance is:\n",
    "{account_details}\n",
    "You have the tools to perform a websearch for relevant news and information.\n",
    "You have tools to check stock prices.\n",
    "You have tools to buy and sell shares.\n",
    "You have tools to save memory of companies, research and thinking so far.\n",
    "Please make use of these tools to manage your portfolio. Carry out trades as you see fit; do not wait for instructions or ask for confirmation.\n",
    "\"\"\"\n",
    "\n",
    "prompt = \"\"\"\n",
    "Use your tools to make decisions about your portfolio.\n",
    "Investigate the news and the market, make your decision, make the trades, and respond with a summary of your actions.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### And to run our Trader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for server in mcp_servers:\n",
    "    await server.connect()\n",
    "\n",
    "researcher_tool = await get_researcher_tool(researcher_mcp_servers)\n",
    "trader_agent = Agent(\n",
    "    name=agent_name,\n",
    "    instructions=instructions,\n",
    "    tools=[researcher_tool],\n",
    "    mcp_servers=trader_mcp_servers,\n",
    "    model=model,\n",
    ")\n",
    "with trace(agent_name):\n",
    "    result = await Runner.run(trader_agent, max_turns=30)\n",
    "display(Markdown(result.final_output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And let's look at the results of the trading\n",
    "await read_accounts_resource(agent_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reviewing the Python module made from this:\n",
    "\n",
    "`mcp_params.py` is where the MCP servers are specified. You'll notice I've brought in some familiar friends: memory and push notifications!\n",
    "\n",
    "`templates.py` is where the instructions and messages are set up (i.e. the System prompts and User prompts)\n",
    "\n",
    "`traders.py` brings it all together.\n",
    "\n",
    "You'll notice I've done something a bit fancy with code like this:\n",
    "\n",
    "```\n",
    "async with AsyncExitStack() as stack:\n",
    "    mcp_servers = [await stack.enter_async_context(MCPServerStdio(params)) for params in mcp_server_params]\n",
    "```\n",
    "\n",
    "This is just a tidy way to combine our \"with\" statements (known as context managers) so that we don't need to do something ugly like this:\n",
    "\n",
    "```\n",
    "async with MCPServerStdio(params=params1) as mcp_server1:\n",
    "    async with MCPServerStdio(params=params2) as mcp_server2:\n",
    "        async with MCPServerStdio(params=params3) as mcp_server3:\n",
    "            mcp_servers = [mcp_server1, mcp_server2, mcp_server3]\n",
    "```\n",
    "\n",
    "But it's equivalent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from traders import Trader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trader = Trader(\"Ian\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await trader.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await read_accounts_resource(\"Ian\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Traces\n",
    "\n",
    "https://platform.openai.com/traces\n",
    "\n",
    "#### Counting number of tools used in total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mcp_params import trader_mcp_server_params, researcher_mcp_server_params\n",
    "\n",
    "all_params = trader_mcp_server_params + researcher_mcp_server_params(\"ed\")\n",
    "\n",
    "count = 0\n",
    "for each_params in all_params:\n",
    "    async with MCPServerStdio(params=each_params, client_session_timeout_seconds=60) as server:\n",
    "        mcp_tools = await server.list_tools()\n",
    "        count += len(mcp_tools)\n",
    "print(f\"We have {len(all_params)} MCP servers, and {count} tools\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
